lib/simd/src/autotune.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 pub const CostDistribution = struct {
  4     num_values: usize = 0,
  5     pointer_padding: [8 - @sizeOf(usize)]u8 = @splat(0),
  6     online_n: f64 = 0,
  7     values: [max_values]f64 = @splat(0),
  8 
  9     pub const max_values: usize = 14;
 10 
 11     const m1_index: usize = 0;
 12     const m2_index: usize = 1;
 13     const mean_index: usize = 2;
 14     const standard_deviation_index: usize = 3;
 15     const lower_index: usize = 4;
 16     const upper_index: usize = 5;
 17 
 18     pub fn notify(self: *@This(), value: f64) void {
 19         if (value < 0) return;
 20         if (self.isOnline()) {
 21             self.onlineNotify(value);
 22             return;
 23         }
 24         std.debug.assert(self.num_values < max_values);
 25         self.values[self.num_values] = value;
 26         self.num_values += 1;
 27         if (self.num_values == max_values) self.warmUpOnline();
 28     }
 29 
 30     pub fn estimateCost(self: *@This()) f64 {
 31         std.debug.assert(self.num_values != 0);
 32         if (!self.isOnline()) self.warmUpOnline();
 33         return self.mean();
 34     }
 35 
 36     pub fn isOnline(self: *const @This()) bool {
 37         return self.online_n > 0;
 38     }
 39 
 40     pub fn bufferedCount(self: *const @This()) usize {
 41         return self.num_values;
 42     }
 43 
 44     pub fn onlineCount(self: *const @This()) f64 {
 45         return self.online_n;
 46     }
 47 
 48     pub fn m1(self: *const @This()) f64 {
 49         return self.values[m1_index];
 50     }
 51 
 52     pub fn m2(self: *const @This()) f64 {
 53         return self.values[m2_index];
 54     }
 55 
 56     pub fn mean(self: *const @This()) f64 {
 57         return self.values[mean_index];
 58     }
 59 
 60     pub fn standardDeviation(self: *const @This()) f64 {
 61         return self.values[standard_deviation_index];
 62     }
 63 
 64     pub fn lower(self: *const @This()) f64 {
 65         return self.values[lower_index];
 66     }
 67 
 68     pub fn upper(self: *const @This()) f64 {
 69         return self.values[upper_index];
 70     }
 71 
 72     fn median(to_sort: []f64) f64 {
 73         std.debug.assert(to_sort.len >= 2);
 74         std.mem.sort(f64, to_sort, {}, std.sort.asc(f64));
 75         const middle = to_sort.len / 2;
 76         if (to_sort.len % 2 != 0) return to_sort[middle];
 77         return (to_sort[middle] + to_sort[middle - 1]) * 0.5;
 78     }
 79 
 80     fn medianAbsoluteDeviation(values: []const f64, median_value: f64) f64 {
 81         var deviations: [max_values]f64 = undefined;
 82         for (values, deviations[0..values.len]) |value, *deviation| {
 83             deviation.* = @abs(value - median_value);
 84         }
 85         return median(deviations[0..values.len]);
 86     }
 87 
 88     fn removeOutliers(self: *@This()) void {
 89         if (self.num_values < 3) return;
 90         std.debug.assert(self.num_values <= max_values);
 91         const active = self.values[0..self.num_values];
 92         const median_value = median(active);
 93         const mad = medianAbsoluteDeviation(active, median_value);
 94         if (mad == 0) {
 95             const skewness = (active[active.len - 1] - median_value) -
 96                 (median_value - active[0]);
 97             const trim = @max(self.num_values / 2, 2);
 98             const left = @max(if (skewness < 0) trim * 3 / 4 else trim / 4, 1);
 99             self.num_values -= trim;
100             std.debug.assert(self.num_values >= 1);
101             std.mem.copyForwards(
102                 f64,
103                 self.values[0..self.num_values],
104                 self.values[left .. left + self.num_values],
105             );
106             return;
107         }
108 
109         const upper_bound = median_value + 5 * mad;
110         const lower_bound = median_value - 5 * mad;
111         var right = self.num_values - 1;
112         while (self.values[right] > upper_bound) right -= 1;
113         std.debug.assert(right >= self.num_values / 2);
114         var left: usize = 0;
115         while (left < right and self.values[left] < lower_bound) left += 1;
116         std.debug.assert(left <= self.num_values / 2);
117         self.num_values = right - left + 1;
118         std.mem.copyForwards(
119             f64,
120             self.values[0..self.num_values],
121             self.values[left .. left + self.num_values],
122         );
123     }
124 
125     fn sampleMean(self: *const @This()) f64 {
126         std.debug.assert(!self.isOnline());
127         std.debug.assert(self.num_values != 0);
128         std.debug.assert(self.num_values <= max_values);
129         var sum: f64 = 0;
130         for (self.values[0..self.num_values]) |value| sum += value;
131         return sum / @as(f64, @floatFromInt(self.num_values));
132     }
133 
134     fn sampleVariance(self: *const @This(), sample_mean: f64) f64 {
135         std.debug.assert(sample_mean >= 0);
136         std.debug.assert(!self.isOnline());
137         std.debug.assert(self.num_values != 0);
138         std.debug.assert(self.num_values <= max_values);
139         if (self.num_values == 1) return 0;
140         var squared_sum: f64 = 0;
141         for (self.values[0..self.num_values]) |value| {
142             const difference = value - sample_mean;
143             squared_sum += difference * difference;
144         }
145         return squared_sum / @as(f64, @floatFromInt(self.num_values - 1));
146     }
147 
148     fn onlineNotify(self: *@This(), unbounded_value: f64) void {
149         const value = @min(@max(self.lower(), unbounded_value), self.upper());
150         const n_minus_one = self.online_n;
151         self.online_n += 1;
152         const difference = value - self.values[m1_index];
153         const difference_div_n = difference / self.online_n;
154         self.values[m1_index] += difference_div_n;
155         std.debug.assert(self.values[m1_index] >= self.lower());
156         self.values[m2_index] += difference * n_minus_one * difference_div_n;
157         const standard_deviation = @sqrt(self.values[m2_index] / @max(1, n_minus_one));
158         self.values[mean_index] = self.values[m1_index] * 0.2 + self.values[mean_index] * 0.8;
159         self.values[standard_deviation_index] = standard_deviation * 0.2 +
160             self.values[standard_deviation_index] * 0.8;
161         self.values[lower_index] = self.values[mean_index] -
162             3.5 * self.values[standard_deviation_index];
163         self.values[upper_index] = self.values[mean_index] +
164             3.5 * self.values[standard_deviation_index];
165     }
166 
167     fn warmUpOnline(self: *@This()) void {
168         self.removeOutliers();
169         const sample_mean = self.sampleMean();
170         const sample_variance = self.sampleVariance(sample_mean);
171         var copy: [max_values]f64 = undefined;
172         @memcpy(copy[0..self.num_values], self.values[0..self.num_values]);
173         self.values[m1_index] = 0;
174         self.values[m2_index] = 0;
175         self.values[mean_index] = sample_mean;
176         self.values[standard_deviation_index] = @sqrt(sample_variance);
177         if (self.values[standard_deviation_index] == 0) {
178             self.values[standard_deviation_index] = self.values[mean_index] / 2;
179         }
180         self.values[lower_index] = self.values[mean_index] -
181             4 * self.values[standard_deviation_index];
182         self.values[upper_index] = self.values[mean_index] +
183             4 * self.values[standard_deviation_index];
184         for (copy[0..self.num_values]) |value| self.onlineNotify(value);
185         std.debug.assert(self.isOnline());
186     }
187 };
188 
189 comptime {
190     if (@sizeOf(CostDistribution) != 128) @compileError("CostDistribution must occupy 128 bytes");
191 }
192 
193 pub fn NextWithSkip(comptime capacity: usize) type {
194     comptime {
195         if (capacity == 0) @compileError("NextWithSkip requires positive capacity");
196         if (capacity >= Link.max_count) {
197             @compileError("NextWithSkip capacity exceeds packed index range");
198         }
199     }
200 
201     return struct {
202         links: [capacity]Link = @splat(.{}),
203         count: usize = 0,
204 
205         const Self = @This();
206 
207         pub fn init(count: usize) Self {
208             std.debug.assert(count != 0);
209             std.debug.assert(count <= capacity);
210             var self = Self{};
211             self.count = count;
212             for (self.links[0..count], 0..) |*link, position| {
213                 link.* = Link.init(position, count);
214             }
215             return self;
216         }
217 
218         pub fn next(self: *const Self, position: usize) usize {
219             std.debug.assert(position < self.count);
220             std.debug.assert(!self.links[position].isRemoved());
221             return self.links[position].next();
222         }
223 
224         pub fn skip(self: *Self, position: usize) void {
225             std.debug.assert(position < self.count);
226             std.debug.assert(!self.links[position].isRemoved());
227             const previous = self.links[position].previous();
228             const next_position = self.links[position].next();
229             if (previous == position or next_position == position) return;
230             self.links[next_position].setPrevious(previous);
231             self.links[previous].setNext(next_position);
232             self.links[position].remove();
233         }
234     };
235 }
236 
237 pub fn AutoTune(
238     comptime Config: type,
239     comptime max_candidates: usize,
240     comptime min_samples: usize,
241 ) type {
242     comptime {
243         if (min_samples < 2) @compileError("AutoTune requires at least two samples");
244     }
245     const CandidateList = NextWithSkip(max_candidates);
246 
247     return struct {
248         best_index: ?usize = null,
249         candidate_values: [max_candidates]Config = undefined,
250         cost_values: [max_candidates]CostDistribution = @splat(.{}),
251         candidate_count: usize = 0,
252         config_index: usize = 0,
253         list: CandidateList = .{},
254         rounds_complete: usize = 0,
255         skip_if_above: f64 = 0,
256 
257         const Self = @This();
258         pub const SetCandidatesError = error{
259             NoCandidates,
260             TooManyCandidates,
261             AlreadyConfigured,
262         };
263 
264         pub fn best(self: *const Self) ?*const Config {
265             const index = self.best_index orelse return null;
266             return &self.candidate_values[index];
267         }
268 
269         pub fn hasCandidates(self: *const Self) bool {
270             return self.candidate_count != 0;
271         }
272 
273         pub fn setCandidates(
274             self: *Self,
275             candidates_to_copy: []const Config,
276         ) SetCandidatesError!void {
277             if (self.hasCandidates()) return error.AlreadyConfigured;
278             if (candidates_to_copy.len == 0) return error.NoCandidates;
279             if (candidates_to_copy.len > max_candidates) return error.TooManyCandidates;
280             @memcpy(self.candidate_values[0..candidates_to_copy.len], candidates_to_copy);
281             for (self.cost_values[0..candidates_to_copy.len]) |*cost| cost.* = .{};
282             self.best_index = null;
283             self.candidate_count = candidates_to_copy.len;
284             self.config_index = 0;
285             self.list = CandidateList.init(candidates_to_copy.len);
286             self.rounds_complete = 0;
287             self.skip_if_above = 0;
288         }
289 
290         pub fn candidates(self: *const Self) []const Config {
291             std.debug.assert(self.hasCandidates());
292             return self.candidate_values[0..self.candidate_count];
293         }
294 
295         pub fn costs(self: *Self) []CostDistribution {
296             return self.cost_values[0..self.candidate_count];
297         }
298 
299         pub fn costsConst(self: *const Self) []const CostDistribution {
300             return self.cost_values[0..self.candidate_count];
301         }
302 
303         pub fn nextConfig(self: *const Self) *const Config {
304             std.debug.assert(self.hasCandidates());
305             return &self.candidate_values[self.config_index];
306         }
307 
308         pub fn currentIndex(self: *const Self) usize {
309             std.debug.assert(self.hasCandidates());
310             return self.config_index;
311         }
312 
313         pub fn completedRounds(self: *const Self) usize {
314             return self.rounds_complete;
315         }
316 
317         pub fn skipThreshold(self: *const Self) f64 {
318             return self.skip_if_above;
319         }
320 
321         pub fn notifyCost(self: *Self, cost: u64) void {
322             std.debug.assert(self.best() == null);
323             std.debug.assert(self.hasCandidates());
324             self.cost_values[self.config_index].notify(@floatFromInt(cost));
325             const measured_index = self.config_index;
326             const measured_cost = if (self.rounds_complete >= min_samples)
327                 self.cost_values[self.config_index].estimateCost()
328             else
329                 0;
330             self.config_index = self.list.next(self.config_index);
331             if (measured_cost > self.skip_if_above) self.list.skip(measured_index);
332             if (self.config_index <= measured_index) {
333                 self.rounds_complete += 1;
334                 if (self.rounds_complete >= min_samples) {
335                     var best_cost = std.math.inf(f64);
336                     var minimum_index: usize = 0;
337                     for (self.cost_values[0..self.candidate_count], 0..) |*distribution, index| {
338                         const estimate = distribution.estimateCost();
339                         if (estimate < best_cost) {
340                             best_cost = estimate;
341                             minimum_index = index;
342                         }
343                     }
344                     self.skip_if_above = best_cost * 1.25;
345                     if (self.rounds_complete == 3 * min_samples / 2 + 1) {
346                         self.best_index = minimum_index;
347                     }
348                 }
349             }
350         }
351 
352         pub fn shouldPrint(self: *const Self) bool {
353             return self.rounds_complete > min_samples;
354         }
355     };
356 }
357 
358 const Link = struct {
359     bits: u32 = 0,
360 
361     const index_bits: u5 = 14;
362     const shift: u5 = 18;
363     const max_count: u32 = 1 << index_bits;
364     const removed_mask: u32 = max_count;
365     const previous_mask: u32 = max_count - 1;
366     const next_clear_mask: u32 = std.math.maxInt(u32) >> index_bits;
367 
368     fn init(position: usize, count: usize) Link {
369         std.debug.assert(count < max_count);
370         const previous_position = if (position == 0) count - 1 else position - 1;
371         const next_position = if (position == count - 1) 0 else position + 1;
372         const link = Link{ .bits = @as(u32, @intCast(next_position)) << shift |
373             @as(u32, @intCast(previous_position)) };
374         std.debug.assert(link.next() == next_position);
375         std.debug.assert(link.previous() == previous_position);
376         std.debug.assert(!link.isRemoved());
377         return link;
378     }
379 
380     fn isRemoved(self: Link) bool {
381         return self.bits & removed_mask != 0;
382     }
383 
384     fn remove(self: *Link) void {
385         self.bits |= removed_mask;
386     }
387 
388     fn next(self: Link) usize {
389         return self.bits >> shift;
390     }
391 
392     fn previous(self: Link) usize {
393         return self.bits & previous_mask;
394     }
395 
396     fn setNext(self: *Link, next_position: usize) void {
397         std.debug.assert(next_position < max_count);
398         self.bits &= next_clear_mask;
399         self.bits |= @as(u32, @intCast(next_position)) << shift;
400         std.debug.assert(self.next() == next_position);
401         std.debug.assert(!self.isRemoved());
402     }
403 
404     fn setPrevious(self: *Link, previous_position: usize) void {
405         std.debug.assert(previous_position < max_count);
406         self.bits &= ~previous_mask;
407         self.bits |= @intCast(previous_position);
408         std.debug.assert(self.previous() == previous_position);
409         std.debug.assert(!self.isRemoved());
410     }
411 };
412 
413 test "Highway cost distribution rejects negative costs and preserves fixed footprint" {
414     try std.testing.expectEqual(@as(usize, 128), @sizeOf(CostDistribution));
415     var distribution = CostDistribution{};
416     distribution.notify(-1);
417     try std.testing.expectEqual(@as(usize, 0), distribution.bufferedCount());
418     distribution.notify(6.5);
419     try std.testing.expectEqual(@as(f64, 6.5), distribution.estimateCost());
420     try std.testing.expect(distribution.isOnline());
421     try std.testing.expectEqual(@as(f64, 1), distribution.onlineCount());
422 }
423 
424 test "Highway cost distribution mitigates equal-sample outliers" {
425     const counts = [_]usize{
426         3,
427         CostDistribution.max_values - 1,
428         CostDistribution.max_values,
429         CostDistribution.max_values + 1,
430     };
431     for (counts) |count| {
432         for ([_]f64{ 0, 1000 }) |outlier| {
433             var distribution = CostDistribution{};
434             const outlier_count = @max(count / 4, 1);
435             for (0..count - outlier_count) |_| distribution.notify(6.5);
436             for (0..outlier_count) |_| distribution.notify(outlier);
437             const estimate = distribution.estimateCost();
438             try std.testing.expect(estimate >= 6.25 and estimate <= 6.75);
439         }
440     }
441 }
442 
443 test "Highway cost distribution matches pinned C++ online states" {
444     var simple = CostDistribution{};
445     for ([_]f64{ 4, 5, 6 }) |value| simple.notify(value);
446     try std.testing.expectEqual(@as(f64, 0x1.32b020c49ba5ep+2), simple.estimateCost());
447     try expectDistribution(&simple, .{
448         0x1.4000000000000p+2,
449         0x1.0000000000000p+1,
450         0x1.32b020c49ba5ep+2,
451         0x1.a6785e357f100p-1,
452         0x1.e76dde34d03b8p+0,
453         0x1.eb84c9fc033cep+2,
454     });
455 
456     var online = CostDistribution{};
457     for ([_]f64{ 0, 9, 10, 11, 12, 10, 9, 1000, 11, 10, 12, 8, 10, 9 }) |value| {
458         online.notify(value);
459     }
460     try expectDistribution(&online, .{
461         0x1.42aaaaaaaaaaap+3,
462         0x1.0eaaaaaaaaaadp+4,
463         0x1.33124f7940cb8p+3,
464         0x1.e85ebad74a944p-1,
465         0x1.907b2d3450f62p+2,
466         0x1.9de70858591bfp+3,
467     });
468     for ([_]f64{ 13, 5000, 7, 10 }) |value| online.notify(value);
469     try std.testing.expectEqual(@as(f64, 0x1.410fef266a414p+3), online.estimateCost());
470     try expectDistribution(&online, .{
471         0x1.48aea7aedd6e6p+3,
472         0x1.684f03f534899p+5,
473         0x1.410fef266a414p+3,
474         0x1.60085b324134cp+0,
475         0x1.4e188e80db746p+2,
476         0x1.db13970c66c85p+3,
477     });
478 }
479 
480 test "Highway cost distribution stabilizes additive and multiplicative noise" {
481     var generator = std.Random.DefaultPrng.init(0x4857_5941_5554_4f54);
482     const random = generator.random();
483     const powers = [_]f64{ 0, 1e3, 1e4, 1e5 };
484     const multipliers = [_]f64{ 0.5, 0.75, 0.85, 0.9 };
485     for (0..100) |_| {
486         var distribution = CostDistribution{};
487         for (0..1000) |_| {
488             var value: f64 = 500;
489             for (0..100) |_| {
490                 const bits: i32 = @intCast(random.int(u32) & 1023);
491                 value += @as(f64, @floatFromInt(bits - 512)) / 64;
492             }
493             const noise = random.int(u32);
494             if (noise < 1 << 28) {
495                 if (noise & 3 != 0) {
496                     value += powers[noise & 3];
497                 } else {
498                     value *= multipliers[noise >> 2 & 3];
499                 }
500             }
501             distribution.notify(value);
502         }
503         const estimate = distribution.estimateCost();
504         try std.testing.expect(estimate >= 490 and estimate <= 540);
505     }
506 }
507 
508 test "Highway next-with-skip wraps and splices edges" {
509     var list = NextWithSkip(123).init(123);
510     try std.testing.expectEqual(@as(usize, 0), list.next(122));
511     try std.testing.expectEqual(@as(usize, 1), list.next(0));
512     list.skip(1);
513     try std.testing.expectEqual(@as(usize, 2), list.next(0));
514     list.skip(2);
515     try std.testing.expectEqual(@as(usize, 3), list.next(0));
516     list.skip(122);
517     try std.testing.expectEqual(@as(usize, 0), list.next(121));
518     list.skip(0);
519     try std.testing.expectEqual(@as(usize, 3), list.next(121));
520 }
521 
522 test "Highway next-with-skip retains the final position" {
523     const counts = [_]usize{ 37, 63, 513 };
524     for (counts) |count| {
525         var list = NextWithSkip(513).init(count);
526         var positions: [513]usize = undefined;
527         for (positions[0..count], 0..) |*position, index| position.* = index;
528         var random = std.Random.DefaultPrng.init(129 * count);
529         random.random().shuffle(usize, positions[0..count]);
530         for (positions[0 .. count - 1]) |position| list.skip(position);
531         const survivor = positions[count - 1];
532         try std.testing.expectEqual(survivor, list.next(survivor));
533         list.skip(survivor);
534         try std.testing.expectEqual(survivor, list.next(survivor));
535     }
536 }
537 
538 test "Highway auto-tune exhaustively selects the first minimum" {
539     const Tuner = AutoTune(u16, 4, 2);
540     var tuner = Tuner{};
541     try std.testing.expect(!tuner.hasCandidates());
542     try std.testing.expectError(error.NoCandidates, tuner.setCandidates(&.{}));
543     try std.testing.expectError(error.TooManyCandidates, tuner.setCandidates(&.{ 0, 1, 2, 3, 4 }));
544     try tuner.setCandidates(&.{ 0, 1, 2, 3 });
545     try std.testing.expectError(error.AlreadyConfigured, tuner.setCandidates(&.{1}));
546     try std.testing.expectEqualSlices(u16, &.{ 0, 1, 2, 3 }, tuner.candidates());
547 
548     const fixed_costs = [_]u64{ 100, 40, 40, 200 };
549     var measurements: [4]usize = @splat(0);
550     var calls: usize = 0;
551     while (tuner.best() == null) : (calls += 1) {
552         try std.testing.expect(calls < 32);
553         const index = tuner.currentIndex();
554         try std.testing.expectEqual(@as(u16, @intCast(index)), tuner.nextConfig().*);
555         measurements[index] += 1;
556         tuner.notifyCost(fixed_costs[index]);
557     }
558     try std.testing.expectEqual(@as(u16, 1), tuner.best().?.*);
559     try std.testing.expectEqual(@as(usize, 4), tuner.completedRounds());
560     try std.testing.expect(tuner.shouldPrint());
561     try std.testing.expectEqualSlices(usize, &.{ 3, 4, 4, 3 }, &measurements);
562     try std.testing.expectEqual(@as(f64, 50), tuner.skipThreshold());
563     try std.testing.expectEqual(@as(usize, 4), tuner.costsConst().len);
564 }
565 
566 fn expectDistribution(
567     distribution: *const CostDistribution,
568     expected: [6]f64,
569 ) !void {
570     try std.testing.expectEqual(expected[0], distribution.m1());
571     try std.testing.expectEqual(expected[1], distribution.m2());
572     try std.testing.expectEqual(expected[2], distribution.mean());
573     try std.testing.expectEqual(expected[3], distribution.standardDeviation());
574     try std.testing.expectEqual(expected[4], distribution.lower());
575     try std.testing.expectEqual(expected[5], distribution.upper());
576 }