tiny.simd.unroll

1pub fn unroll(unit: anytype, input: anytype, output: anytype) void {2    const Unit = pointee(@TypeOf(unit));3    const Base = unitBase(Unit, 1);4    if (!@hasDecl(Unit, "func")) @compileError("unroller unit must declare func");5    const input_slice: []const Base.Input = input;6    const output_slice: []Base.Output = output;7    const lane_count = Base.actual_lanes;8    std.debug.assert(input_slice.len <= std.math.maxInt(isize));910    var x0 = x0Init(Unit, Base, unit);11    var y = yInit(Unit, Base, unit);12    var index: usize = 0;1314    if (input_slice.len < lane_count) {15        const count: isize = @intCast(input_slice.len);16        x0 = maskLoad(Unit, Base, unit, 0, input_slice, count);17        y = unit.func(0, x0, y);18        _ = maskStore(Unit, Base, unit, 0, output_slice, y, count);19        _ = reduceFinal(Unit, Base, unit, y, output_slice);20        return;21    }2223    if (input_slice.len > 4 * lane_count) {24        var x01 = x0Init(Unit, Base, unit);25        var y1 = yInit(Unit, Base, unit);26        var x02 = x0Init(Unit, Base, unit);27        var y2 = yInit(Unit, Base, unit);28        var x03 = x0Init(Unit, Base, unit);29        var y3 = yInit(Unit, Base, unit);3031        while (index + 4 * lane_count <= input_slice.len) {32            x0 = load(Unit, Base, unit, @intCast(index), input_slice);33            x01 = load(Unit, Base, unit, @intCast(index + lane_count), input_slice);34            x02 = load(Unit, Base, unit, @intCast(index + 2 * lane_count), input_slice);35            x03 = load(Unit, Base, unit, @intCast(index + 3 * lane_count), input_slice);3637            y = unit.func(@intCast(index), x0, y);38            y1 = unit.func(@intCast(index + lane_count), x01, y1);39            y2 = unit.func(@intCast(index + 2 * lane_count), x02, y2);40            y3 = unit.func(@intCast(index + 3 * lane_count), x03, y3);4142            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;43            index += lane_count;44            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y1)) return;45            index += lane_count;46            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y2)) return;47            index += lane_count;48            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y3)) return;49            index += lane_count;50        }5152        y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);53    }5455    while (index + lane_count <= input_slice.len) : (index += lane_count) {56        x0 = load(Unit, Base, unit, @intCast(index), input_slice);57        y = unit.func(@intCast(index), x0, y);58        if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;59    }6061    if (index != input_slice.len) {62        const tail_index = input_slice.len - lane_count;63        const places = @as(isize, @intCast(index)) - @as(isize, @intCast(input_slice.len));64        x0 = maskLoad(Unit, Base, unit, @intCast(tail_index), input_slice, places);65        y = unit.func(@intCast(tail_index), x0, y);66        _ = maskStore(Unit, Base, unit, @intCast(tail_index), output_slice, y, places);67    }6869    _ = reduceFinal(Unit, Base, unit, y, output_slice);70}

1pub const unroll = unroller.unroll;