tiny.simd.unroller

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

1const std = @import("std");2const memory = @import("memory.zig");3const tag = @import("tag.zig");45const Window = struct {6    lane_start: usize,7    storage_start: usize,8    count: usize,9};1011pub 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);1617    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;3132        pub fn maxUnitLanes() usize {33            return max_unit_lanes;34        }3536        pub fn actualLanes() usize {37            return actual_lanes;38        }3940        pub fn x0Init() InputVector {41            return @splat(0);42        }4344        pub fn yInit() OutputVector {45            return @splat(0);46        }4748        pub fn load(index: isize, input: []const Input) InputVector {49            return loadFull(InputTag, index, input);50        }5152        pub fn maskLoad(index: isize, input: []const Input, places: isize) InputVector {53            return maskLoadOr(@splat(0), index, input, places);54        }5556        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        }6465        pub fn storeAndShortCircuit(66            index: isize,67            output: []Output,68            value: OutputVector,69        ) bool {70            storeFull(OutputTag, index, output, value);71            return true;72        }7374        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        }8283        pub fn reduceFinal(_: OutputVector, _: []Output) usize {84            return 0;85        }8687        pub fn reduceUnrolled(88            _: OutputVector,89            _: OutputVector,90            _: OutputVector,91            value: OutputVector,92        ) OutputVector {93            return value;94        }95    };96}9798pub 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);111112    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;130131        pub fn maxUnitLanes() usize {132            return max_unit_lanes;133        }134135        pub fn actualLanes() usize {136            return actual_lanes;137        }138139        pub fn x0Init() Input0Vector {140            return @splat(0);141        }142143        pub fn x1Init() Input1Vector {144            return @splat(0);145        }146147        pub fn yInit() OutputVector {148            return @splat(0);149        }150151        pub fn load0(index: isize, input: []const Input0) Input0Vector {152            return loadFull(Input0Tag, index, input);153        }154155        pub fn load1(index: isize, input: []const Input1) Input1Vector {156            return loadFull(Input1Tag, index, input);157        }158159        pub fn maskLoad0(160            index: isize,161            input: []const Input0,162            places: isize,163        ) Input0Vector {164            return maskLoad0Or(@splat(0), index, input, places);165        }166167        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        }175176        pub fn maskLoad1(177            index: isize,178            input: []const Input1,179            places: isize,180        ) Input1Vector {181            return maskLoad1Or(@splat(0), index, input, places);182        }183184        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        }192193        pub fn storeAndShortCircuit(194            index: isize,195            output: []Output,196            value: OutputVector,197        ) bool {198            storeFull(OutputTag, index, output, value);199            return true;200        }201202        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        }210211        pub fn reduceFinal(_: OutputVector, _: []Output) usize {212            return 0;213        }214215        pub fn reduceUnrolled(216            _: OutputVector,217            _: OutputVector,218            _: OutputVector,219            value: OutputVector,220        ) OutputVector {221            return value;222        }223    };224}225226pub 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));234235    var x0 = x0Init(Unit, Base, unit);236    var y = yInit(Unit, Base, unit);237    var index: usize = 0;238239    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    }247248    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);255256        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);261262            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);266267            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        }276277        y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);278    }279280    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    }285286    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    }293294    _ = reduceFinal(Unit, Base, unit, y, output_slice);295}296297pub 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));307308    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;312313    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    }322323    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);333334        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);343344            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);348349            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        }358359        y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);360    }361362    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    }368369    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    }377378    _ = reduceFinal(Unit, Base, unit, y, output_slice);379}380381fn 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}385386fn 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}390391fn 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}395396fn 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}400401fn 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}411412fn 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}422423fn 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}433434fn 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}445446fn 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}457458fn 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}469470fn 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}483484fn 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}498499fn 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}509510fn 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}524525fn 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}532533fn 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}547548fn 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}560561fn 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}575576fn 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}595596fn pointee(comptime Pointer: type) type {597    return switch (@typeInfo(Pointer)) {598        .pointer => |info| if (info.size == .one)599            info.child600        else601            @compileError("unroller unit must be passed by single-item pointer"),602        else => @compileError("unroller unit must be passed by pointer"),603    };604}605606fn 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}616617fn validateLane(comptime T: type) void {618    if (!tag.isLane(T)) @compileError("unroller units require SIMD lane types");619}620621fn 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}630631fn MultiplyUnit(comptime T: type) type {632    const UnitBase = UnrollerUnit2D(T, T, T);633    return struct {634        pub const Base: type = UnitBase;635636        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}647648fn AccumulateUnit(comptime T: type) type {649    const UnitBase = UnrollerUnit(T, T);650    return struct {651        pub const Base: type = UnitBase;652653        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        }661662        pub fn storeAndShortCircuit(663            _: *@This(),664            _: isize,665            _: []T,666            _: Base.OutputVector,667        ) bool {668            return true;669        }670671        pub fn maskStore(672            _: *@This(),673            _: isize,674            _: []T,675            _: Base.OutputVector,676            _: isize,677        ) usize {678            return 0;679        }680681        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        }690691        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}702703fn DotUnit(comptime T: type) type {704    const UnitBase = UnrollerUnit2D(T, T, T);705    return struct {706        pub const Base: type = UnitBase;707708        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        }717718        pub fn storeAndShortCircuit(719            _: *@This(),720            _: isize,721            _: []T,722            _: Base.OutputVector,723        ) bool {724            return true;725        }726727        pub fn maskStore(728            _: *@This(),729            _: isize,730            _: []T,731            _: Base.OutputVector,732            _: isize,733        ) usize {734            return 0;735        }736737        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        }746747        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}758759fn MinUnit(comptime T: type) type {760    const UnitBase = UnrollerUnit(T, T);761    return struct {762        pub const Base: type = UnitBase;763764        pub fn func(765            _: *@This(),766            _: isize,767            x: Base.InputVector,768            y: Base.OutputVector,769        ) Base.OutputVector {770            return @min(x, y);771        }772773        pub fn yInit(_: *@This()) Base.OutputVector {774            return @splat(highestValue(T));775        }776777        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        }785786        pub fn storeAndShortCircuit(787            _: *@This(),788            _: isize,789            _: []T,790            _: Base.OutputVector,791        ) bool {792            return true;793        }794795        pub fn maskStore(796            _: *@This(),797            _: isize,798            _: []T,799            _: Base.OutputVector,800            _: isize,801        ) usize {802            return 0;803        }804805        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        }814815        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}826827fn ConvertUnit(comptime From: type, comptime To: type) type {828    const UnitBase = UnrollerUnit(From, To);829    return struct {830        pub const Base: type = UnitBase;831832        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}846847fn FindUnit(comptime T: type) type {848    const Index = signedLane(@sizeOf(T));849    const UnitBase = UnrollerUnit(T, Index);850    return struct {851        needle: T,852853        pub const Base: type = UnitBase;854855        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        }869870        pub fn x0Init(self: *@This()) Base.InputVector {871            return @splat(otherValue(T, self.needle));872        }873874        pub fn yInit(_: *@This()) Base.OutputVector {875            return @splat(-1);876        }877878        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        }886887        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        }897898        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}911912const 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,923924    pub const Base: type = UnrollerUnit(i32, i32);925926    pub fn x0Init(self: *@This()) Base.InputVector {927        self.x0_inits += 1;928        return Base.x0Init();929    }930931    pub fn yInit(self: *@This()) Base.OutputVector {932        self.y_inits += 1;933        return Base.yInit();934    }935936    pub fn load(self: *@This(), index: isize, input: []const i32) Base.InputVector {937        self.loads += 1;938        return Base.load(index, input);939    }940941    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    }951952    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    }961962    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    }971972    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    }983984    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    }994995    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};10041005const StopUnit = struct {1006    funcs: usize = 0,1007    stores: usize = 0,1008    final_reductions: usize = 0,10091010    pub const Base: type = UnrollerUnit(i32, i32);10111012    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    }10211022    pub fn storeAndShortCircuit(1023        self: *@This(),1024        _: isize,1025        _: []i32,1026        _: Base.OutputVector,1027    ) bool {1028        self.stores += 1;1029        return false;1030    }10311032    pub fn reduceFinal(1033        self: *@This(),1034        _: Base.OutputVector,1035        _: []i32,1036    ) usize {1037        self.final_reductions += 1;1038        return 0;1039    }1040};10411042fn 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}10601061fn 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);10811082        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);10871088        var dot = DotUnit(T){};1089        var direct: [1]T = undefined;1090        unroll2(&dot, a[0..count], b[0..count], &direct);10911092        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);11001101        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}11071108fn 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}11301131fn 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;11391140        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])]);11441145        var zero = Unit{ .needle = 0 };1146        unroll(&zero, input[0..count], &found);1147        try std.testing.expectEqual(@as(Index, 0), found[0]);11481149        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}11541155fn addVectors(comptime T: type, a: anytype, b: @TypeOf(a)) @TypeOf(a) {1156    return if (comptime @typeInfo(T) == .int) a +% b else a + b;1157}11581159fn multiplyVectors(comptime T: type, a: anytype, b: @TypeOf(a)) @TypeOf(a) {1160    return if (comptime @typeInfo(T) == .int) a *% b else a * b;1161}11621163fn 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        else1169            result + value[index];1170    }1171    return result;1172}11731174fn 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}11811182fn 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}11891190fn 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}11971198fn 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}12131214test "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}12261227test "Highway unroller ports multiply dot sum and minimum sweeps" {1228    inline for (.{ f16, f32, f64 }) |T| try verifyDotAndMin(T);1229}12301231test "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}12371238test "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}12611262test "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]);12751276    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);12891290    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}