tiny.simd.autotune

1const std = @import("std");23pub 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),89    pub const max_values: usize = 14;1011    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;1718    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    }2930    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    }3536    pub fn isOnline(self: *const @This()) bool {37        return self.online_n > 0;38    }3940    pub fn bufferedCount(self: *const @This()) usize {41        return self.num_values;42    }4344    pub fn onlineCount(self: *const @This()) f64 {45        return self.online_n;46    }4748    pub fn m1(self: *const @This()) f64 {49        return self.values[m1_index];50    }5152    pub fn m2(self: *const @This()) f64 {53        return self.values[m2_index];54    }5556    pub fn mean(self: *const @This()) f64 {57        return self.values[mean_index];58    }5960    pub fn standardDeviation(self: *const @This()) f64 {61        return self.values[standard_deviation_index];62    }6364    pub fn lower(self: *const @This()) f64 {65        return self.values[lower_index];66    }6768    pub fn upper(self: *const @This()) f64 {69        return self.values[upper_index];70    }7172    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    }7980    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    }8788    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        }108109        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    }124125    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    }133134    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    }147148    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    }166167    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};188189comptime {190    if (@sizeOf(CostDistribution) != 128) @compileError("CostDistribution must occupy 128 bytes");191}192193pub 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    }200201    return struct {202        links: [capacity]Link = @splat(.{}),203        count: usize = 0,204205        const Self = @This();206207        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        }217218        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        }223224        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}236237pub 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);246247    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,256257        const Self = @This();258        pub const SetCandidatesError = error{259            NoCandidates,260            TooManyCandidates,261            AlreadyConfigured,262        };263264        pub fn best(self: *const Self) ?*const Config {265            const index = self.best_index orelse return null;266            return &self.candidate_values[index];267        }268269        pub fn hasCandidates(self: *const Self) bool {270            return self.candidate_count != 0;271        }272273        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        }289290        pub fn candidates(self: *const Self) []const Config {291            std.debug.assert(self.hasCandidates());292            return self.candidate_values[0..self.candidate_count];293        }294295        pub fn costs(self: *Self) []CostDistribution {296            return self.cost_values[0..self.candidate_count];297        }298299        pub fn costsConst(self: *const Self) []const CostDistribution {300            return self.cost_values[0..self.candidate_count];301        }302303        pub fn nextConfig(self: *const Self) *const Config {304            std.debug.assert(self.hasCandidates());305            return &self.candidate_values[self.config_index];306        }307308        pub fn currentIndex(self: *const Self) usize {309            std.debug.assert(self.hasCandidates());310            return self.config_index;311        }312313        pub fn completedRounds(self: *const Self) usize {314            return self.rounds_complete;315        }316317        pub fn skipThreshold(self: *const Self) f64 {318            return self.skip_if_above;319        }320321        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            else329                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        }351352        pub fn shouldPrint(self: *const Self) bool {353            return self.rounds_complete > min_samples;354        }355    };356}357358const Link = struct {359    bits: u32 = 0,360361    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;367368    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    }379380    fn isRemoved(self: Link) bool {381        return self.bits & removed_mask != 0;382    }383384    fn remove(self: *Link) void {385        self.bits |= removed_mask;386    }387388    fn next(self: Link) usize {389        return self.bits >> shift;390    }391392    fn previous(self: Link) usize {393        return self.bits & previous_mask;394    }395396    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    }403404    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};412413test "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}423424test "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}442443test "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    });455456    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}479480test "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}507508test "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}521522test "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}537538test "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());547548    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}565566fn 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}

1pub const autotune = @import("autotune.zig");