lib/simd/src/unroller.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const memory = @import("memory.zig");
3 const tag = @import("tag.zig");
4
5 const Window = struct {
6 lane_start: usize,
7 storage_start: usize,
8 count: usize,
9 };
10
11 pub fn UnrollerUnit(comptime InputLane: type, comptime OutputLane: type) type {
12 validateLane(InputLane);
13 validateLane(OutputLane);
14 const LargerLane = signedLane(@max(@sizeOf(InputLane), @sizeOf(OutputLane)));
15 const LargerTag = tag.ScalableTag(LargerLane);
16
17 return struct {
18 pub const input_count: usize = 1;
19 pub const Input: type = InputLane;
20 pub const Output: type = OutputLane;
21 pub const Larger: type = LargerLane;
22 pub const Tag: type = LargerTag;
23 pub const InputTag: type = LargerTag.rebind(InputLane);
24 pub const OutputTag: type = LargerTag.rebind(OutputLane);
25 pub const InputVector: type = InputTag.Vector;
26 pub const OutputVector: type = OutputTag.Vector;
27 pub const XVector: type = InputVector;
28 pub const YVector: type = OutputVector;
29 pub const max_unit_lanes: usize = LargerTag.lane_count;
30 pub const actual_lanes: usize = LargerTag.lane_count;
31
32 pub fn maxUnitLanes() usize {
33 return max_unit_lanes;
34 }
35
36 pub fn actualLanes() usize {
37 return actual_lanes;
38 }
39
40 pub fn x0Init() InputVector {
41 return @splat(0);
42 }
43
44 pub fn yInit() OutputVector {
45 return @splat(0);
46 }
47
48 pub fn load(index: isize, input: []const Input) InputVector {
49 return loadFull(InputTag, index, input);
50 }
51
52 pub fn maskLoad(index: isize, input: []const Input, places: isize) InputVector {
53 return maskLoadOr(@splat(0), index, input, places);
54 }
55
56 pub fn maskLoadOr(
57 inactive: InputVector,
58 index: isize,
59 input: []const Input,
60 places: isize,
61 ) InputVector {
62 return loadPartial(InputTag, inactive, index, input, places);
63 }
64
65 pub fn storeAndShortCircuit(
66 index: isize,
67 output: []Output,
68 value: OutputVector,
69 ) bool {
70 storeFull(OutputTag, index, output, value);
71 return true;
72 }
73
74 pub fn maskStore(
75 index: isize,
76 output: []Output,
77 value: OutputVector,
78 places: isize,
79 ) usize {
80 return storePartial(OutputTag, index, output, value, places);
81 }
82
83 pub fn reduceFinal(_: OutputVector, _: []Output) usize {
84 return 0;
85 }
86
87 pub fn reduceUnrolled(
88 _: OutputVector,
89 _: OutputVector,
90 _: OutputVector,
91 value: OutputVector,
92 ) OutputVector {
93 return value;
94 }
95 };
96 }
97
98 pub fn UnrollerUnit2D(
99 comptime Input0Lane: type,
100 comptime Input1Lane: type,
101 comptime OutputLane: type,
102 ) type {
103 validateLane(Input0Lane);
104 validateLane(Input1Lane);
105 validateLane(OutputLane);
106 const LargerLane = signedLane(@max(
107 @sizeOf(Input0Lane),
108 @max(@sizeOf(Input1Lane), @sizeOf(OutputLane)),
109 ));
110 const LargerTag = tag.ScalableTag(LargerLane);
111
112 return struct {
113 pub const input_count: usize = 2;
114 pub const Input0: type = Input0Lane;
115 pub const Input1: type = Input1Lane;
116 pub const Output: type = OutputLane;
117 pub const Larger: type = LargerLane;
118 pub const Tag: type = LargerTag;
119 pub const Input0Tag: type = LargerTag.rebind(Input0Lane);
120 pub const Input1Tag: type = LargerTag.rebind(Input1Lane);
121 pub const OutputTag: type = LargerTag.rebind(OutputLane);
122 pub const Input0Vector: type = Input0Tag.Vector;
123 pub const Input1Vector: type = Input1Tag.Vector;
124 pub const OutputVector: type = OutputTag.Vector;
125 pub const X0Vector: type = Input0Vector;
126 pub const X1Vector: type = Input1Vector;
127 pub const YVector: type = OutputVector;
128 pub const max_unit_lanes: usize = LargerTag.lane_count;
129 pub const actual_lanes: usize = LargerTag.lane_count;
130
131 pub fn maxUnitLanes() usize {
132 return max_unit_lanes;
133 }
134
135 pub fn actualLanes() usize {
136 return actual_lanes;
137 }
138
139 pub fn x0Init() Input0Vector {
140 return @splat(0);
141 }
142
143 pub fn x1Init() Input1Vector {
144 return @splat(0);
145 }
146
147 pub fn yInit() OutputVector {
148 return @splat(0);
149 }
150
151 pub fn load0(index: isize, input: []const Input0) Input0Vector {
152 return loadFull(Input0Tag, index, input);
153 }
154
155 pub fn load1(index: isize, input: []const Input1) Input1Vector {
156 return loadFull(Input1Tag, index, input);
157 }
158
159 pub fn maskLoad0(
160 index: isize,
161 input: []const Input0,
162 places: isize,
163 ) Input0Vector {
164 return maskLoad0Or(@splat(0), index, input, places);
165 }
166
167 pub fn maskLoad0Or(
168 inactive: Input0Vector,
169 index: isize,
170 input: []const Input0,
171 places: isize,
172 ) Input0Vector {
173 return loadPartial(Input0Tag, inactive, index, input, places);
174 }
175
176 pub fn maskLoad1(
177 index: isize,
178 input: []const Input1,
179 places: isize,
180 ) Input1Vector {
181 return maskLoad1Or(@splat(0), index, input, places);
182 }
183
184 pub fn maskLoad1Or(
185 inactive: Input1Vector,
186 index: isize,
187 input: []const Input1,
188 places: isize,
189 ) Input1Vector {
190 return loadPartial(Input1Tag, inactive, index, input, places);
191 }
192
193 pub fn storeAndShortCircuit(
194 index: isize,
195 output: []Output,
196 value: OutputVector,
197 ) bool {
198 storeFull(OutputTag, index, output, value);
199 return true;
200 }
201
202 pub fn maskStore(
203 index: isize,
204 output: []Output,
205 value: OutputVector,
206 places: isize,
207 ) usize {
208 return storePartial(OutputTag, index, output, value, places);
209 }
210
211 pub fn reduceFinal(_: OutputVector, _: []Output) usize {
212 return 0;
213 }
214
215 pub fn reduceUnrolled(
216 _: OutputVector,
217 _: OutputVector,
218 _: OutputVector,
219 value: OutputVector,
220 ) OutputVector {
221 return value;
222 }
223 };
224 }
225
226 pub fn unroll(unit: anytype, input: anytype, output: anytype) void {
227 const Unit = pointee(@TypeOf(unit));
228 const Base = unitBase(Unit, 1);
229 if (!@hasDecl(Unit, "func")) @compileError("unroller unit must declare func");
230 const input_slice: []const Base.Input = input;
231 const output_slice: []Base.Output = output;
232 const lane_count = Base.actual_lanes;
233 std.debug.assert(input_slice.len <= std.math.maxInt(isize));
234
235 var x0 = x0Init(Unit, Base, unit);
236 var y = yInit(Unit, Base, unit);
237 var index: usize = 0;
238
239 if (input_slice.len < lane_count) {
240 const count: isize = @intCast(input_slice.len);
241 x0 = maskLoad(Unit, Base, unit, 0, input_slice, count);
242 y = unit.func(0, x0, y);
243 _ = maskStore(Unit, Base, unit, 0, output_slice, y, count);
244 _ = reduceFinal(Unit, Base, unit, y, output_slice);
245 return;
246 }
247
248 if (input_slice.len > 4 * lane_count) {
249 var x01 = x0Init(Unit, Base, unit);
250 var y1 = yInit(Unit, Base, unit);
251 var x02 = x0Init(Unit, Base, unit);
252 var y2 = yInit(Unit, Base, unit);
253 var x03 = x0Init(Unit, Base, unit);
254 var y3 = yInit(Unit, Base, unit);
255
256 while (index + 4 * lane_count <= input_slice.len) {
257 x0 = load(Unit, Base, unit, @intCast(index), input_slice);
258 x01 = load(Unit, Base, unit, @intCast(index + lane_count), input_slice);
259 x02 = load(Unit, Base, unit, @intCast(index + 2 * lane_count), input_slice);
260 x03 = load(Unit, Base, unit, @intCast(index + 3 * lane_count), input_slice);
261
262 y = unit.func(@intCast(index), x0, y);
263 y1 = unit.func(@intCast(index + lane_count), x01, y1);
264 y2 = unit.func(@intCast(index + 2 * lane_count), x02, y2);
265 y3 = unit.func(@intCast(index + 3 * lane_count), x03, y3);
266
267 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;
268 index += lane_count;
269 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y1)) return;
270 index += lane_count;
271 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y2)) return;
272 index += lane_count;
273 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y3)) return;
274 index += lane_count;
275 }
276
277 y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);
278 }
279
280 while (index + lane_count <= input_slice.len) : (index += lane_count) {
281 x0 = load(Unit, Base, unit, @intCast(index), input_slice);
282 y = unit.func(@intCast(index), x0, y);
283 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;
284 }
285
286 if (index != input_slice.len) {
287 const tail_index = input_slice.len - lane_count;
288 const places = @as(isize, @intCast(index)) - @as(isize, @intCast(input_slice.len));
289 x0 = maskLoad(Unit, Base, unit, @intCast(tail_index), input_slice, places);
290 y = unit.func(@intCast(tail_index), x0, y);
291 _ = maskStore(Unit, Base, unit, @intCast(tail_index), output_slice, y, places);
292 }
293
294 _ = reduceFinal(Unit, Base, unit, y, output_slice);
295 }
296
297 pub fn unroll2(unit: anytype, input0: anytype, input1: anytype, output: anytype) void {
298 const Unit = pointee(@TypeOf(unit));
299 const Base = unitBase(Unit, 2);
300 if (!@hasDecl(Unit, "func")) @compileError("two-input unroller unit must declare func");
301 const input0_slice: []const Base.Input0 = input0;
302 const input1_slice: []const Base.Input1 = input1;
303 const output_slice: []Base.Output = output;
304 const lane_count = Base.actual_lanes;
305 std.debug.assert(input0_slice.len == input1_slice.len);
306 std.debug.assert(input0_slice.len <= std.math.maxInt(isize));
307
308 var x0 = x0Init2(Unit, Base, unit);
309 var x1 = x1Init2(Unit, Base, unit);
310 var y = yInit(Unit, Base, unit);
311 var index: usize = 0;
312
313 if (input0_slice.len < lane_count) {
314 const count: isize = @intCast(input0_slice.len);
315 x0 = maskLoad0(Unit, Base, unit, 0, input0_slice, count);
316 x1 = maskLoad1(Unit, Base, unit, 0, input1_slice, count);
317 y = unit.func(0, x0, x1, y);
318 _ = maskStore(Unit, Base, unit, 0, output_slice, y, count);
319 _ = reduceFinal(Unit, Base, unit, y, output_slice);
320 return;
321 }
322
323 if (input0_slice.len > 4 * lane_count) {
324 var x01 = x0Init2(Unit, Base, unit);
325 var x11 = x1Init2(Unit, Base, unit);
326 var y1 = yInit(Unit, Base, unit);
327 var x02 = x0Init2(Unit, Base, unit);
328 var x12 = x1Init2(Unit, Base, unit);
329 var y2 = yInit(Unit, Base, unit);
330 var x03 = x0Init2(Unit, Base, unit);
331 var x13 = x1Init2(Unit, Base, unit);
332 var y3 = yInit(Unit, Base, unit);
333
334 while (index + 4 * lane_count <= input0_slice.len) {
335 x0 = load0(Unit, Base, unit, @intCast(index), input0_slice);
336 x1 = load1(Unit, Base, unit, @intCast(index), input1_slice);
337 x01 = load0(Unit, Base, unit, @intCast(index + lane_count), input0_slice);
338 x11 = load1(Unit, Base, unit, @intCast(index + lane_count), input1_slice);
339 x02 = load0(Unit, Base, unit, @intCast(index + 2 * lane_count), input0_slice);
340 x12 = load1(Unit, Base, unit, @intCast(index + 2 * lane_count), input1_slice);
341 x03 = load0(Unit, Base, unit, @intCast(index + 3 * lane_count), input0_slice);
342 x13 = load1(Unit, Base, unit, @intCast(index + 3 * lane_count), input1_slice);
343
344 y = unit.func(@intCast(index), x0, x1, y);
345 y1 = unit.func(@intCast(index + lane_count), x01, x11, y1);
346 y2 = unit.func(@intCast(index + 2 * lane_count), x02, x12, y2);
347 y3 = unit.func(@intCast(index + 3 * lane_count), x03, x13, y3);
348
349 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;
350 index += lane_count;
351 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y1)) return;
352 index += lane_count;
353 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y2)) return;
354 index += lane_count;
355 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y3)) return;
356 index += lane_count;
357 }
358
359 y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);
360 }
361
362 while (index + lane_count <= input0_slice.len) : (index += lane_count) {
363 x0 = load0(Unit, Base, unit, @intCast(index), input0_slice);
364 x1 = load1(Unit, Base, unit, @intCast(index), input1_slice);
365 y = unit.func(@intCast(index), x0, x1, y);
366 if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;
367 }
368
369 if (index != input0_slice.len) {
370 const tail_index = input0_slice.len - lane_count;
371 const places = @as(isize, @intCast(index)) - @as(isize, @intCast(input0_slice.len));
372 x0 = maskLoad0(Unit, Base, unit, @intCast(tail_index), input0_slice, places);
373 x1 = maskLoad1(Unit, Base, unit, @intCast(tail_index), input1_slice, places);
374 y = unit.func(@intCast(tail_index), x0, x1, y);
375 _ = maskStore(Unit, Base, unit, @intCast(tail_index), output_slice, y, places);
376 }
377
378 _ = reduceFinal(Unit, Base, unit, y, output_slice);
379 }
380
381 fn x0Init(comptime Unit: type, comptime Base: type, unit: anytype) Base.InputVector {
382 if (comptime @hasDecl(Unit, "x0Init")) return unit.x0Init();
383 return Base.x0Init();
384 }
385
386 fn x0Init2(comptime Unit: type, comptime Base: type, unit: anytype) Base.Input0Vector {
387 if (comptime @hasDecl(Unit, "x0Init")) return unit.x0Init();
388 return Base.x0Init();
389 }
390
391 fn x1Init2(comptime Unit: type, comptime Base: type, unit: anytype) Base.Input1Vector {
392 if (comptime @hasDecl(Unit, "x1Init")) return unit.x1Init();
393 return Base.x1Init();
394 }
395
396 fn yInit(comptime Unit: type, comptime Base: type, unit: anytype) Base.OutputVector {
397 if (comptime @hasDecl(Unit, "yInit")) return unit.yInit();
398 return Base.yInit();
399 }
400
401 fn load(
402 comptime Unit: type,
403 comptime Base: type,
404 unit: anytype,
405 index: isize,
406 input: []const Base.Input,
407 ) Base.InputVector {
408 if (comptime @hasDecl(Unit, "load")) return unit.load(index, input);
409 return Base.load(index, input);
410 }
411
412 fn load0(
413 comptime Unit: type,
414 comptime Base: type,
415 unit: anytype,
416 index: isize,
417 input: []const Base.Input0,
418 ) Base.Input0Vector {
419 if (comptime @hasDecl(Unit, "load0")) return unit.load0(index, input);
420 return Base.load0(index, input);
421 }
422
423 fn load1(
424 comptime Unit: type,
425 comptime Base: type,
426 unit: anytype,
427 index: isize,
428 input: []const Base.Input1,
429 ) Base.Input1Vector {
430 if (comptime @hasDecl(Unit, "load1")) return unit.load1(index, input);
431 return Base.load1(index, input);
432 }
433
434 fn maskLoad(
435 comptime Unit: type,
436 comptime Base: type,
437 unit: anytype,
438 index: isize,
439 input: []const Base.Input,
440 places: isize,
441 ) Base.InputVector {
442 if (comptime @hasDecl(Unit, "maskLoad")) return unit.maskLoad(index, input, places);
443 return Base.maskLoad(index, input, places);
444 }
445
446 fn maskLoad0(
447 comptime Unit: type,
448 comptime Base: type,
449 unit: anytype,
450 index: isize,
451 input: []const Base.Input0,
452 places: isize,
453 ) Base.Input0Vector {
454 if (comptime @hasDecl(Unit, "maskLoad0")) return unit.maskLoad0(index, input, places);
455 return Base.maskLoad0(index, input, places);
456 }
457
458 fn maskLoad1(
459 comptime Unit: type,
460 comptime Base: type,
461 unit: anytype,
462 index: isize,
463 input: []const Base.Input1,
464 places: isize,
465 ) Base.Input1Vector {
466 if (comptime @hasDecl(Unit, "maskLoad1")) return unit.maskLoad1(index, input, places);
467 return Base.maskLoad1(index, input, places);
468 }
469
470 fn storeAndShortCircuit(
471 comptime Unit: type,
472 comptime Base: type,
473 unit: anytype,
474 index: isize,
475 output: []Base.Output,
476 value: Base.OutputVector,
477 ) bool {
478 if (comptime @hasDecl(Unit, "storeAndShortCircuit")) {
479 return unit.storeAndShortCircuit(index, output, value);
480 }
481 return Base.storeAndShortCircuit(index, output, value);
482 }
483
484 fn maskStore(
485 comptime Unit: type,
486 comptime Base: type,
487 unit: anytype,
488 index: isize,
489 output: []Base.Output,
490 value: Base.OutputVector,
491 places: isize,
492 ) usize {
493 if (comptime @hasDecl(Unit, "maskStore")) {
494 return unit.maskStore(index, output, value, places);
495 }
496 return Base.maskStore(index, output, value, places);
497 }
498
499 fn reduceFinal(
500 comptime Unit: type,
501 comptime Base: type,
502 unit: anytype,
503 value: Base.OutputVector,
504 output: []Base.Output,
505 ) usize {
506 if (comptime @hasDecl(Unit, "reduceFinal")) return unit.reduceFinal(value, output);
507 return Base.reduceFinal(value, output);
508 }
509
510 fn reduceUnrolled(
511 comptime Unit: type,
512 comptime Base: type,
513 unit: anytype,
514 x0: Base.OutputVector,
515 x1: Base.OutputVector,
516 x2: Base.OutputVector,
517 value: Base.OutputVector,
518 ) Base.OutputVector {
519 if (comptime @hasDecl(Unit, "reduceUnrolled")) {
520 return unit.reduceUnrolled(x0, x1, x2, value);
521 }
522 return Base.reduceUnrolled(x0, x1, x2, value);
523 }
524
525 fn loadFull(comptime D: type, index: isize, input: []const D.Lane) D.Vector {
526 std.debug.assert(index >= 0);
527 const start: usize = @intCast(index);
528 std.debug.assert(start <= input.len);
529 std.debug.assert(D.lane_count <= input.len - start);
530 return memory.load(D, input[start..]);
531 }
532
533 fn loadPartial(
534 comptime D: type,
535 inactive: D.Vector,
536 index: isize,
537 input: []const D.Lane,
538 places: isize,
539 ) D.Vector {
540 const selected = window(D.lane_count, index, places, input.len);
541 var lanes: [D.lane_count]D.Lane = inactive;
542 for (0..selected.count) |offset| {
543 lanes[selected.lane_start + offset] = input[selected.storage_start + offset];
544 }
545 return lanes;
546 }
547
548 fn storeFull(
549 comptime D: type,
550 index: isize,
551 output: []D.Lane,
552 value: D.Vector,
553 ) void {
554 std.debug.assert(index >= 0);
555 const start: usize = @intCast(index);
556 std.debug.assert(start <= output.len);
557 std.debug.assert(D.lane_count <= output.len - start);
558 memory.store(D, value, output[start..]);
559 }
560
561 fn storePartial(
562 comptime D: type,
563 index: isize,
564 output: []D.Lane,
565 value: D.Vector,
566 places: isize,
567 ) usize {
568 const selected = window(D.lane_count, index, places, output.len);
569 const lanes: [D.lane_count]D.Lane = value;
570 for (0..selected.count) |offset| {
571 output[selected.storage_start + offset] = lanes[selected.lane_start + offset];
572 }
573 return selected.count;
574 }
575
576 fn window(lane_count: usize, index: isize, places: isize, length: usize) Window {
577 const bound: isize = @intCast(lane_count);
578 std.debug.assert(places >= -bound);
579 std.debug.assert(places <= bound);
580 const count: usize = @intCast(if (places < 0) -places else places);
581 const lane_start = if (places < 0) lane_count - count else 0;
582 const lane_offset: isize = @intCast(lane_start);
583 std.debug.assert(index <= std.math.maxInt(isize) - lane_offset);
584 const storage_index = index + lane_offset;
585 std.debug.assert(storage_index >= 0);
586 const storage_start: usize = @intCast(storage_index);
587 std.debug.assert(storage_start <= length);
588 std.debug.assert(count <= length - storage_start);
589 return .{
590 .lane_start = lane_start,
591 .storage_start = storage_start,
592 .count = count,
593 };
594 }
595
596 fn pointee(comptime Pointer: type) type {
597 return switch (@typeInfo(Pointer)) {
598 .pointer => |info| if (info.size == .one)
599 info.child
600 else
601 @compileError("unroller unit must be passed by single-item pointer"),
602 else => @compileError("unroller unit must be passed by pointer"),
603 };
604 }
605
606 fn unitBase(comptime Unit: type, comptime input_count: usize) type {
607 if (!@hasDecl(Unit, "Base")) {
608 @compileError("unroller unit must declare Base = UnrollerUnit(...) or UnrollerUnit2D(...)");
609 }
610 const Base = Unit.Base;
611 if (!@hasDecl(Base, "input_count") or Base.input_count != input_count) {
612 @compileError("unroller unit Base has the wrong input count");
613 }
614 return Base;
615 }
616
617 fn validateLane(comptime T: type) void {
618 if (!tag.isLane(T)) @compileError("unroller units require SIMD lane types");
619 }
620
621 fn signedLane(comptime byte_count: usize) type {
622 return switch (byte_count) {
623 1 => i8,
624 2 => i16,
625 4 => i32,
626 8 => i64,
627 else => @compileError("unroller lane sizes must be 1, 2, 4, or 8 bytes"),
628 };
629 }
630
631 fn MultiplyUnit(comptime T: type) type {
632 const UnitBase = UnrollerUnit2D(T, T, T);
633 return struct {
634 pub const Base: type = UnitBase;
635
636 pub fn func(
637 _: *@This(),
638 _: isize,
639 x0: Base.Input0Vector,
640 x1: Base.Input1Vector,
641 _: Base.OutputVector,
642 ) Base.OutputVector {
643 return multiplyVectors(T, x0, x1);
644 }
645 };
646 }
647
648 fn AccumulateUnit(comptime T: type) type {
649 const UnitBase = UnrollerUnit(T, T);
650 return struct {
651 pub const Base: type = UnitBase;
652
653 pub fn func(
654 _: *@This(),
655 _: isize,
656 x: Base.InputVector,
657 y: Base.OutputVector,
658 ) Base.OutputVector {
659 return addVectors(T, x, y);
660 }
661
662 pub fn storeAndShortCircuit(
663 _: *@This(),
664 _: isize,
665 _: []T,
666 _: Base.OutputVector,
667 ) bool {
668 return true;
669 }
670
671 pub fn maskStore(
672 _: *@This(),
673 _: isize,
674 _: []T,
675 _: Base.OutputVector,
676 _: isize,
677 ) usize {
678 return 0;
679 }
680
681 pub fn reduceFinal(
682 _: *@This(),
683 value: Base.OutputVector,
684 output: []T,
685 ) usize {
686 std.debug.assert(output.len >= 1);
687 output[0] = sumVector(T, value);
688 return 1;
689 }
690
691 pub fn reduceUnrolled(
692 _: *@This(),
693 x0: Base.OutputVector,
694 x1: Base.OutputVector,
695 x2: Base.OutputVector,
696 value: Base.OutputVector,
697 ) Base.OutputVector {
698 return addVectors(T, addVectors(T, x0, x1), addVectors(T, x2, value));
699 }
700 };
701 }
702
703 fn DotUnit(comptime T: type) type {
704 const UnitBase = UnrollerUnit2D(T, T, T);
705 return struct {
706 pub const Base: type = UnitBase;
707
708 pub fn func(
709 _: *@This(),
710 _: isize,
711 x0: Base.Input0Vector,
712 x1: Base.Input1Vector,
713 y: Base.OutputVector,
714 ) Base.OutputVector {
715 return addVectors(T, multiplyVectors(T, x0, x1), y);
716 }
717
718 pub fn storeAndShortCircuit(
719 _: *@This(),
720 _: isize,
721 _: []T,
722 _: Base.OutputVector,
723 ) bool {
724 return true;
725 }
726
727 pub fn maskStore(
728 _: *@This(),
729 _: isize,
730 _: []T,
731 _: Base.OutputVector,
732 _: isize,
733 ) usize {
734 return 0;
735 }
736
737 pub fn reduceFinal(
738 _: *@This(),
739 value: Base.OutputVector,
740 output: []T,
741 ) usize {
742 std.debug.assert(output.len >= 1);
743 output[0] = sumVector(T, value);
744 return 1;
745 }
746
747 pub fn reduceUnrolled(
748 _: *@This(),
749 x0: Base.OutputVector,
750 x1: Base.OutputVector,
751 x2: Base.OutputVector,
752 value: Base.OutputVector,
753 ) Base.OutputVector {
754 return addVectors(T, addVectors(T, x0, x1), addVectors(T, x2, value));
755 }
756 };
757 }
758
759 fn MinUnit(comptime T: type) type {
760 const UnitBase = UnrollerUnit(T, T);
761 return struct {
762 pub const Base: type = UnitBase;
763
764 pub fn func(
765 _: *@This(),
766 _: isize,
767 x: Base.InputVector,
768 y: Base.OutputVector,
769 ) Base.OutputVector {
770 return @min(x, y);
771 }
772
773 pub fn yInit(_: *@This()) Base.OutputVector {
774 return @splat(highestValue(T));
775 }
776
777 pub fn maskLoad(
778 _: *@This(),
779 index: isize,
780 input: []const T,
781 places: isize,
782 ) Base.InputVector {
783 return Base.maskLoadOr(@splat(highestValue(T)), index, input, places);
784 }
785
786 pub fn storeAndShortCircuit(
787 _: *@This(),
788 _: isize,
789 _: []T,
790 _: Base.OutputVector,
791 ) bool {
792 return true;
793 }
794
795 pub fn maskStore(
796 _: *@This(),
797 _: isize,
798 _: []T,
799 _: Base.OutputVector,
800 _: isize,
801 ) usize {
802 return 0;
803 }
804
805 pub fn reduceFinal(
806 _: *@This(),
807 value: Base.OutputVector,
808 output: []T,
809 ) usize {
810 std.debug.assert(output.len >= 1);
811 output[0] = minVector(T, value);
812 return 1;
813 }
814
815 pub fn reduceUnrolled(
816 _: *@This(),
817 x0: Base.OutputVector,
818 x1: Base.OutputVector,
819 x2: Base.OutputVector,
820 value: Base.OutputVector,
821 ) Base.OutputVector {
822 return @min(@min(x0, x1), @min(x2, value));
823 }
824 };
825 }
826
827 fn ConvertUnit(comptime From: type, comptime To: type) type {
828 const UnitBase = UnrollerUnit(From, To);
829 return struct {
830 pub const Base: type = UnitBase;
831
832 pub fn func(
833 _: *@This(),
834 _: isize,
835 x: Base.InputVector,
836 _: Base.OutputVector,
837 ) Base.OutputVector {
838 var result: Base.OutputVector = undefined;
839 inline for (0..Base.actual_lanes) |lane_index| {
840 result[lane_index] = convertScalar(To, x[lane_index]);
841 }
842 return result;
843 }
844 };
845 }
846
847 fn FindUnit(comptime T: type) type {
848 const Index = signedLane(@sizeOf(T));
849 const UnitBase = UnrollerUnit(T, Index);
850 return struct {
851 needle: T,
852
853 pub const Base: type = UnitBase;
854
855 pub fn func(
856 self: *@This(),
857 index: isize,
858 x: Base.InputVector,
859 y: Base.OutputVector,
860 ) Base.OutputVector {
861 inline for (0..Base.actual_lanes) |lane_index| {
862 if (x[lane_index] == self.needle) {
863 const found = index + @as(isize, @intCast(lane_index));
864 return @splat(@as(Index, @intCast(found)));
865 }
866 }
867 return y;
868 }
869
870 pub fn x0Init(self: *@This()) Base.InputVector {
871 return @splat(otherValue(T, self.needle));
872 }
873
874 pub fn yInit(_: *@This()) Base.OutputVector {
875 return @splat(-1);
876 }
877
878 pub fn maskLoad(
879 self: *@This(),
880 index: isize,
881 input: []const T,
882 places: isize,
883 ) Base.InputVector {
884 return Base.maskLoadOr(@splat(otherValue(T, self.needle)), index, input, places);
885 }
886
887 pub fn storeAndShortCircuit(
888 _: *@This(),
889 _: isize,
890 output: []Index,
891 value: Base.OutputVector,
892 ) bool {
893 std.debug.assert(output.len >= 1);
894 output[0] = value[0];
895 return value[0] == -1;
896 }
897
898 pub fn maskStore(
899 _: *@This(),
900 _: isize,
901 output: []Index,
902 value: Base.OutputVector,
903 _: isize,
904 ) usize {
905 std.debug.assert(output.len >= 1);
906 output[0] = value[0];
907 return 1;
908 }
909 };
910 }
911
912 const TraceUnit = struct {
913 x0_inits: usize = 0,
914 y_inits: usize = 0,
915 loads: usize = 0,
916 mask_loads: usize = 0,
917 funcs: usize = 0,
918 stores: usize = 0,
919 mask_stores: usize = 0,
920 unrolled_reductions: usize = 0,
921 final_reductions: usize = 0,
922 last_places: isize = 0,
923
924 pub const Base: type = UnrollerUnit(i32, i32);
925
926 pub fn x0Init(self: *@This()) Base.InputVector {
927 self.x0_inits += 1;
928 return Base.x0Init();
929 }
930
931 pub fn yInit(self: *@This()) Base.OutputVector {
932 self.y_inits += 1;
933 return Base.yInit();
934 }
935
936 pub fn load(self: *@This(), index: isize, input: []const i32) Base.InputVector {
937 self.loads += 1;
938 return Base.load(index, input);
939 }
940
941 pub fn maskLoad(
942 self: *@This(),
943 index: isize,
944 input: []const i32,
945 places: isize,
946 ) Base.InputVector {
947 self.mask_loads += 1;
948 self.last_places = places;
949 return Base.maskLoad(index, input, places);
950 }
951
952 pub fn func(
953 self: *@This(),
954 _: isize,
955 x: Base.InputVector,
956 _: Base.OutputVector,
957 ) Base.OutputVector {
958 self.funcs += 1;
959 return x;
960 }
961
962 pub fn storeAndShortCircuit(
963 self: *@This(),
964 index: isize,
965 output: []i32,
966 value: Base.OutputVector,
967 ) bool {
968 self.stores += 1;
969 return Base.storeAndShortCircuit(index, output, value);
970 }
971
972 pub fn maskStore(
973 self: *@This(),
974 index: isize,
975 output: []i32,
976 value: Base.OutputVector,
977 places: isize,
978 ) usize {
979 self.mask_stores += 1;
980 self.last_places = places;
981 return Base.maskStore(index, output, value, places);
982 }
983
984 pub fn reduceUnrolled(
985 self: *@This(),
986 x0: Base.OutputVector,
987 x1: Base.OutputVector,
988 x2: Base.OutputVector,
989 value: Base.OutputVector,
990 ) Base.OutputVector {
991 self.unrolled_reductions += 1;
992 return Base.reduceUnrolled(x0, x1, x2, value);
993 }
994
995 pub fn reduceFinal(
996 self: *@This(),
997 value: Base.OutputVector,
998 output: []i32,
999 ) usize {
1000 self.final_reductions += 1;
1001 return Base.reduceFinal(value, output);
1002 }
1003 };
1004
1005 const StopUnit = struct {
1006 funcs: usize = 0,
1007 stores: usize = 0,
1008 final_reductions: usize = 0,
1009
1010 pub const Base: type = UnrollerUnit(i32, i32);
1011
1012 pub fn func(
1013 self: *@This(),
1014 _: isize,
1015 x: Base.InputVector,
1016 _: Base.OutputVector,
1017 ) Base.OutputVector {
1018 self.funcs += 1;
1019 return x;
1020 }
1021
1022 pub fn storeAndShortCircuit(
1023 self: *@This(),
1024 _: isize,
1025 _: []i32,
1026 _: Base.OutputVector,
1027 ) bool {
1028 self.stores += 1;
1029 return false;
1030 }
1031
1032 pub fn reduceFinal(
1033 self: *@This(),
1034 _: Base.OutputVector,
1035 _: []i32,
1036 ) usize {
1037 self.final_reductions += 1;
1038 return 0;
1039 }
1040 };
1041
1042 fn counts(lane_count: usize) [14]usize {
1043 return .{
1044 1,
1045 3,
1046 7,
1047 16,
1048 @max(lane_count / 2, 1),
1049 @max(2 * lane_count / 3, 1),
1050 lane_count,
1051 lane_count + 1,
1052 4 * lane_count / 3,
1053 3 * lane_count,
1054 8 * lane_count,
1055 8 * lane_count + 2,
1056 256 * lane_count - 1,
1057 256 * lane_count,
1058 };
1059 }
1060
1061 fn verifyDotAndMin(comptime T: type) !void {
1062 const Base = DotUnit(T).Base;
1063 const max_values = 256 * Base.actual_lanes;
1064 var a: [max_values]T = undefined;
1065 var b: [max_values]T = undefined;
1066 var products: [max_values]T = undefined;
1067 for (counts(Base.actual_lanes)) |count| {
1068 var expected_dot_f64: f64 = 0;
1069 var expected_min = std.math.floatMax(T);
1070 for (a[0..count], b[0..count], 0..) |*a_value, *b_value, index| {
1071 const a_integer = @as(i32, @intCast(index * 37 % 1_024)) - 512;
1072 const b_integer = @as(i32, @intCast(index * 53 % 1_024)) - 512;
1073 a_value.* = @as(T, @floatFromInt(a_integer)) * @as(T, 1.0 / 64.0);
1074 b_value.* = @as(T, @floatFromInt(b_integer)) * @as(T, 1.0 / 64.0);
1075 expected_dot_f64 += @as(f64, @floatCast(a_value.*)) *
1076 @as(f64, @floatCast(b_value.*));
1077 expected_min = @min(expected_min, a_value.*);
1078 }
1079 const expected_dot: T = @floatCast(expected_dot_f64);
1080 expected_dot_f64 = @floatCast(expected_dot);
1081
1082 var multiply = MultiplyUnit(T){};
1083 unroll2(&multiply, a[0..count], b[0..count], products[0..count]);
1084 var accumulate = AccumulateUnit(T){};
1085 var via_multiply: [1]T = undefined;
1086 unroll(&accumulate, products[0..count], &via_multiply);
1087
1088 var dot = DotUnit(T){};
1089 var direct: [1]T = undefined;
1090 unroll2(&dot, a[0..count], b[0..count], &direct);
1091
1092 const tolerance = 120.0 * @as(f64, @floatCast(std.math.floatEps(T))) *
1093 @abs(expected_dot_f64);
1094 const multiply_error = @abs(
1095 expected_dot_f64 - @as(f64, @floatCast(via_multiply[0])),
1096 );
1097 const direct_error = @abs(expected_dot_f64 - @as(f64, @floatCast(direct[0])));
1098 try std.testing.expect(multiply_error <= tolerance);
1099 try std.testing.expect(direct_error <= tolerance);
1100
1101 var minimum = MinUnit(T){};
1102 var actual_min: [1]T = undefined;
1103 unroll(&minimum, a[0..count], &actual_min);
1104 try std.testing.expectEqual(expected_min, actual_min[0]);
1105 }
1106 }
1107
1108 fn verifyConvert(comptime T: type) !void {
1109 const Base = ConvertUnit(T, i32).Base;
1110 const max_values = 256 * Base.actual_lanes;
1111 var input: [max_values]T = undefined;
1112 var integers: [max_values]i32 = undefined;
1113 var round_trip: [max_values]T = undefined;
1114 for (counts(Base.actual_lanes)) |count| {
1115 for (input[0..count], 0..) |*value, index| {
1116 value.* = @as(T, @floatFromInt(index)) * @as(T, 0.25);
1117 }
1118 var to_integer = ConvertUnit(T, i32){};
1119 unroll(&to_integer, input[0..count], integers[0..count]);
1120 for (input[0..count], integers[0..count]) |value, integer| {
1121 try std.testing.expectEqual(@as(i32, @intFromFloat(value)), integer);
1122 }
1123 var to_float = ConvertUnit(i32, T){};
1124 unroll(&to_float, integers[0..count], round_trip[0..count]);
1125 for (integers[0..count], round_trip[0..count]) |integer, value| {
1126 try std.testing.expectEqual(@as(T, @floatFromInt(integer)), value);
1127 }
1128 }
1129 }
1130
1131 fn verifyFind(comptime T: type) !void {
1132 const Unit = FindUnit(T);
1133 const Index = Unit.Base.Output;
1134 const max_values = 256 * Unit.Base.actual_lanes;
1135 var input: [max_values]T = undefined;
1136 for (counts(Unit.Base.actual_lanes)) |count| {
1137 for (input[0..count], 0..) |*value, index| value.* = @floatFromInt(index);
1138 var found: [1]Index = undefined;
1139
1140 var last = Unit{ .needle = @floatFromInt(count - 1) };
1141 unroll(&last, input[0..count], &found);
1142 try std.testing.expect(found[0] >= 0);
1143 try std.testing.expectEqual(last.needle, input[@intCast(found[0])]);
1144
1145 var zero = Unit{ .needle = 0 };
1146 unroll(&zero, input[0..count], &found);
1147 try std.testing.expectEqual(@as(Index, 0), found[0]);
1148
1149 var absent = Unit{ .needle = std.math.floatMax(T) };
1150 unroll(&absent, input[0..count], &found);
1151 try std.testing.expectEqual(@as(Index, -1), found[0]);
1152 }
1153 }
1154
1155 fn addVectors(comptime T: type, a: anytype, b: @TypeOf(a)) @TypeOf(a) {
1156 return if (comptime @typeInfo(T) == .int) a +% b else a + b;
1157 }
1158
1159 fn multiplyVectors(comptime T: type, a: anytype, b: @TypeOf(a)) @TypeOf(a) {
1160 return if (comptime @typeInfo(T) == .int) a *% b else a * b;
1161 }
1162
1163 fn sumVector(comptime T: type, value: anytype) T {
1164 var result: T = 0;
1165 inline for (0..@typeInfo(@TypeOf(value)).vector.len) |index| {
1166 result = if (comptime @typeInfo(T) == .int)
1167 result +% value[index]
1168 else
1169 result + value[index];
1170 }
1171 return result;
1172 }
1173
1174 fn minVector(comptime T: type, value: anytype) T {
1175 var result = value[0];
1176 inline for (1..@typeInfo(@TypeOf(value)).vector.len) |index| {
1177 result = @min(result, value[index]);
1178 }
1179 return result;
1180 }
1181
1182 fn highestValue(comptime T: type) T {
1183 return switch (@typeInfo(T)) {
1184 .int => std.math.maxInt(T),
1185 .float => std.math.floatMax(T),
1186 else => unreachable,
1187 };
1188 }
1189
1190 fn otherValue(comptime T: type, value: T) T {
1191 return switch (@typeInfo(T)) {
1192 .int => value +% 1,
1193 .float => std.math.nan(T),
1194 else => unreachable,
1195 };
1196 }
1197
1198 fn convertScalar(comptime To: type, value: anytype) To {
1199 return switch (@typeInfo(@TypeOf(value))) {
1200 .int => switch (@typeInfo(To)) {
1201 .int => @intCast(value),
1202 .float => @floatFromInt(value),
1203 else => unreachable,
1204 },
1205 .float => switch (@typeInfo(To)) {
1206 .int => @intFromFloat(value),
1207 .float => @floatCast(value),
1208 else => unreachable,
1209 },
1210 else => unreachable,
1211 };
1212 }
1213
1214 test "Highway unroller descriptors use the largest lane size" {
1215 const One = UnrollerUnit(u8, f64);
1216 const Two = UnrollerUnit2D(f16, u32, f64);
1217 try std.testing.expectEqual(tag.ScalableTag(i64).lane_count, One.actual_lanes);
1218 try std.testing.expectEqual(One.actual_lanes, One.InputTag.lane_count);
1219 try std.testing.expectEqual(One.actual_lanes, One.OutputTag.lane_count);
1220 try std.testing.expectEqual(One.max_unit_lanes, One.maxUnitLanes());
1221 try std.testing.expectEqual(tag.ScalableTag(i64).lane_count, Two.actual_lanes);
1222 try std.testing.expectEqual(Two.actual_lanes, Two.Input0Tag.lane_count);
1223 try std.testing.expectEqual(Two.actual_lanes, Two.Input1Tag.lane_count);
1224 try std.testing.expectEqual(Two.actual_lanes, Two.OutputTag.lane_count);
1225 }
1226
1227 test "Highway unroller ports multiply dot sum and minimum sweeps" {
1228 inline for (.{ f16, f32, f64 }) |T| try verifyDotAndMin(T);
1229 }
1230
1231 test "Highway unroller ports widening narrowing and find sweeps" {
1232 inline for (.{ f16, f32, f64 }) |T| {
1233 try verifyConvert(T);
1234 try verifyFind(T);
1235 }
1236 }
1237
1238 test "Highway unroller invokes hooks across four-way and overlapping tails" {
1239 const lanes = TraceUnit.Base.actual_lanes;
1240 const tail_count = @min(2, lanes - 1);
1241 const has_tail: usize = @intFromBool(tail_count != 0);
1242 const count = 5 * lanes + tail_count;
1243 var input: [count]i32 = undefined;
1244 var output: [input.len]i32 = undefined;
1245 for (&input, 0..) |*value, index| value.* = @intCast(index * 7 + 3);
1246 var trace = TraceUnit{};
1247 unroll(&trace, &input, &output);
1248 try std.testing.expectEqualSlices(i32, &input, &output);
1249 try std.testing.expectEqual(@as(usize, 4), trace.x0_inits);
1250 try std.testing.expectEqual(@as(usize, 4), trace.y_inits);
1251 try std.testing.expectEqual(@as(usize, 5), trace.loads);
1252 try std.testing.expectEqual(has_tail, trace.mask_loads);
1253 try std.testing.expectEqual(5 + has_tail, trace.funcs);
1254 try std.testing.expectEqual(@as(usize, 5), trace.stores);
1255 try std.testing.expectEqual(has_tail, trace.mask_stores);
1256 try std.testing.expectEqual(@as(usize, 1), trace.unrolled_reductions);
1257 try std.testing.expectEqual(@as(usize, 1), trace.final_reductions);
1258 try std.testing.expectEqual(-@as(isize, @intCast(tail_count)), trace.last_places);
1259 try std.testing.expectEqual(count, output.len);
1260 }
1261
1262 test "Highway unroller handles small spans and short circuits before reduction" {
1263 var empty_input: [0]i32 = .{};
1264 var empty_output: [0]i32 = .{};
1265 var empty_trace = TraceUnit{};
1266 unroll(&empty_trace, &empty_input, &empty_output);
1267 try std.testing.expectEqual(@as(usize, 1), empty_trace.funcs);
1268 try std.testing.expectEqual(@as(usize, 1), empty_trace.mask_loads);
1269 try std.testing.expectEqual(@as(usize, 1), empty_trace.mask_stores);
1270 try std.testing.expectEqual(@as(usize, 1), empty_trace.final_reductions);
1271 var sum = AccumulateUnit(i32){};
1272 var empty_sum = [_]i32{99};
1273 unroll(&sum, &empty_input, &empty_sum);
1274 try std.testing.expectEqual(@as(i32, 0), empty_sum[0]);
1275
1276 const small_count = @min(3, TraceUnit.Base.actual_lanes - 1);
1277 var small_input: [small_count]i32 = undefined;
1278 var small_output: [small_count]i32 = @splat(0);
1279 for (&small_input, 0..) |*value, index| value.* = @intCast(index + 4);
1280 var trace = TraceUnit{};
1281 unroll(&trace, &small_input, &small_output);
1282 try std.testing.expectEqualSlices(i32, &small_input, &small_output);
1283 try std.testing.expectEqual(@as(usize, 1), trace.x0_inits);
1284 try std.testing.expectEqual(@as(usize, 1), trace.y_inits);
1285 try std.testing.expectEqual(@as(usize, 0), trace.loads);
1286 try std.testing.expectEqual(@as(usize, 1), trace.mask_loads);
1287 try std.testing.expectEqual(@as(usize, 1), trace.mask_stores);
1288 try std.testing.expectEqual(@as(isize, @intCast(small_count)), trace.last_places);
1289
1290 const count = 5 * StopUnit.Base.actual_lanes;
1291 var input: [count]i32 = @splat(1);
1292 var output: [count]i32 = @splat(0);
1293 var stop = StopUnit{};
1294 unroll(&stop, &input, &output);
1295 try std.testing.expectEqual(@as(usize, 4), stop.funcs);
1296 try std.testing.expectEqual(@as(usize, 1), stop.stores);
1297 try std.testing.expectEqual(@as(usize, 0), stop.final_reductions);
1298 }