tiny.simd.unroll2

1pub fn unroll2(unit: anytype, input0: anytype, input1: anytype, output: anytype) void {2    const Unit = pointee(@TypeOf(unit));3    const Base = unitBase(Unit, 2);4    if (!@hasDecl(Unit, "func")) @compileError("two-input unroller unit must declare func");5    const input0_slice: []const Base.Input0 = input0;6    const input1_slice: []const Base.Input1 = input1;7    const output_slice: []Base.Output = output;8    const lane_count = Base.actual_lanes;9    std.debug.assert(input0_slice.len == input1_slice.len);10    std.debug.assert(input0_slice.len <= std.math.maxInt(isize));1112    var x0 = x0Init2(Unit, Base, unit);13    var x1 = x1Init2(Unit, Base, unit);14    var y = yInit(Unit, Base, unit);15    var index: usize = 0;1617    if (input0_slice.len < lane_count) {18        const count: isize = @intCast(input0_slice.len);19        x0 = maskLoad0(Unit, Base, unit, 0, input0_slice, count);20        x1 = maskLoad1(Unit, Base, unit, 0, input1_slice, count);21        y = unit.func(0, x0, x1, y);22        _ = maskStore(Unit, Base, unit, 0, output_slice, y, count);23        _ = reduceFinal(Unit, Base, unit, y, output_slice);24        return;25    }2627    if (input0_slice.len > 4 * lane_count) {28        var x01 = x0Init2(Unit, Base, unit);29        var x11 = x1Init2(Unit, Base, unit);30        var y1 = yInit(Unit, Base, unit);31        var x02 = x0Init2(Unit, Base, unit);32        var x12 = x1Init2(Unit, Base, unit);33        var y2 = yInit(Unit, Base, unit);34        var x03 = x0Init2(Unit, Base, unit);35        var x13 = x1Init2(Unit, Base, unit);36        var y3 = yInit(Unit, Base, unit);3738        while (index + 4 * lane_count <= input0_slice.len) {39            x0 = load0(Unit, Base, unit, @intCast(index), input0_slice);40            x1 = load1(Unit, Base, unit, @intCast(index), input1_slice);41            x01 = load0(Unit, Base, unit, @intCast(index + lane_count), input0_slice);42            x11 = load1(Unit, Base, unit, @intCast(index + lane_count), input1_slice);43            x02 = load0(Unit, Base, unit, @intCast(index + 2 * lane_count), input0_slice);44            x12 = load1(Unit, Base, unit, @intCast(index + 2 * lane_count), input1_slice);45            x03 = load0(Unit, Base, unit, @intCast(index + 3 * lane_count), input0_slice);46            x13 = load1(Unit, Base, unit, @intCast(index + 3 * lane_count), input1_slice);4748            y = unit.func(@intCast(index), x0, x1, y);49            y1 = unit.func(@intCast(index + lane_count), x01, x11, y1);50            y2 = unit.func(@intCast(index + 2 * lane_count), x02, x12, y2);51            y3 = unit.func(@intCast(index + 3 * lane_count), x03, x13, y3);5253            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;54            index += lane_count;55            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y1)) return;56            index += lane_count;57            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y2)) return;58            index += lane_count;59            if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y3)) return;60            index += lane_count;61        }6263        y = reduceUnrolled(Unit, Base, unit, y3, y2, y1, y);64    }6566    while (index + lane_count <= input0_slice.len) : (index += lane_count) {67        x0 = load0(Unit, Base, unit, @intCast(index), input0_slice);68        x1 = load1(Unit, Base, unit, @intCast(index), input1_slice);69        y = unit.func(@intCast(index), x0, x1, y);70        if (!storeAndShortCircuit(Unit, Base, unit, @intCast(index), output_slice, y)) return;71    }7273    if (index != input0_slice.len) {74        const tail_index = input0_slice.len - lane_count;75        const places = @as(isize, @intCast(index)) - @as(isize, @intCast(input0_slice.len));76        x0 = maskLoad0(Unit, Base, unit, @intCast(tail_index), input0_slice, places);77        x1 = maskLoad1(Unit, Base, unit, @intCast(tail_index), input1_slice, places);78        y = unit.func(@intCast(tail_index), x0, x1, y);79        _ = maskStore(Unit, Base, unit, @intCast(tail_index), output_slice, y, places);80    }8182    _ = reduceFinal(Unit, Base, unit, y, output_slice);83}

1pub const unroll2 = unroller.unroll2;